import yaml
import traceback
from abc import ABC, abstractmethod
from typing import Any, Dict, Optional, List

import logging
from copy import deepcopy
from langgraph.types import Send
from langgraph.config import get_stream_writer
from langgraph.types import Command
from langgraph.graph.state import END

from csagent.core.config.schema import NodeConfig, EdgeCondition
from csagent.core.context import AgentState
from csagent.core.utils import evaluate_expression
from csagent.core.constants import END_NODE

logger = logging.getLogger()
class BaseNode(ABC):
    def __init__(self, node_config: NodeConfig):
        self.node_config = node_config
        self.edge_condition: List[EdgeCondition] = []
        self.conf = self._load_node_conf(node_config.conf)
        self.agent_name = ''
        self.initialize()
    
    def initialize(self):
        """初始化节点"""
        pass

    def _load_node_conf(self, config_path: Optional[str]) -> Dict:
        """加载节点配置"""
        if not config_path:
            return {}
        # 实现配置加载逻辑
        with open(config_path, 'r') as file:
            config = yaml.safe_load(file)
            return config
    def _parallel_route(self, dest: str, state: AgentState, batch: int) -> List[Send]:
        num = 0
        one_batch = {}
        route = []
        for id, item in state['parallel_info'].items():
            one_batch[id] = item
            num += 1
            if num >= batch:
                parallel_state = deepcopy(state)
                parallel_state['parallel_info'] = one_batch
                one_batch = {}
                num = 0
                send_dest = Send(dest, parallel_state)
                route.append(send_dest)
        if one_batch:
            parallel_state = deepcopy(state)
            parallel_state['parallel_info'] = one_batch
            send_dest = Send(dest, parallel_state)
            route.append(send_dest)
        return route

    def run(self, state: AgentState) -> Command:
        try:
            stream_writer = get_stream_writer()
            result = self.execute(state, stream_writer)
        except Exception as e:
            if state.get('debug', False):
                trace_info = traceback.format_exc()
            else:
                trace_info = ''
            err_msg = f"agent[{self.agent_name}],{self.node_config.name}节点执行异常, 异常类型: {type(e).__name__}, 消息: {str(e)}, 详细信息: {trace_info}"
            logger.error(err_msg)
            state['code'] = 2000
            state['err_msg'] = err_msg
            
            return Command(
                # this is the state update
                update = state,
                # this is a replacement for an edge
                goto = [END],
            )
        if result['code'] != 0:
            return Command(
                # this is the state update
                update = result,
                # this is a replacement for an edge
                goto = [END],
            )
        goto = []
        default = None
        default_edge = None
        for edge in self.edge_condition:
            dest = edge.name
            if dest == END_NODE:
                dest = END
            route = False
            if not edge.condition:
                route = True
                #goto.append(dest)
            elif edge.condition == "default":
                default = dest
                default_edge = edge
            else:
                flag = False
                try:
                    flag = evaluate_expression(edge.condition, result)
                except Exception as e:
                    if state.get('debug', False):
                        trace_info = traceback.format_exc()
                    else:
                        trace_info = ''
                    err_msg = f"agent[{self.agent_name}],路由condition计算错误, source: {self.node_config.name}, dest:{edge.name}, 异常类型: {type(e).__name__}, 消息: {str(e)},详细信息: {trace_info}"
                    logger.error(err_msg)
                    result['code'] = 2000
                    result['err_msg'] = err_msg
                    flag = False
                    goto = [END]
                    break
                if flag == True:
                    route = True
                    #goto.append(dest)
            if not route:
                continue
            if edge.parallel.get('enable', False) == True:
                batch = int(edge.parallel.get('batch', 1))
                goto.extend(self._parallel_route(dest, result, batch))
            else:
                goto.append(dest)

        if not goto and default:
            if default_edge.parallel.get('enable', False) == True:
                batch = int(default_edge.parallel.get('batch', 1))
                goto.extend(self._parallel_route(default, result, batch))
            else:
                goto.append(default)

        if not goto:
            goto.append(END)
        logger.info(f"agent[{self.agent_name}],route: {self.node_config.name}->{goto}")
        return Command(
            # this is the state update
            update = result,
            # this is a replacement for an edge
            goto = goto,
        )
    
    @abstractmethod
    def execute(self, state: AgentState, stream_writer=None) -> AgentState:
        """
        执行节点逻辑
        :param context: 执行上下文
        :return: 更新后的上下文
        """
        pass
